import torch
import torch.onnx

# 加载保存的模型，然后导出为ONNX格式
# model.load_state_dict(torch.load("/hy-tmp/result/pure-mobilenet.pth"))
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = torch.load('/hy-tmp/result/pure-mobilenet-model.pth')
# 设置模型为评估模式
model.eval()
model = model.to(device)
model = model.eval()
# print("Start inference")
input = torch.randn(1, 3, 224, 224).float().to(device)
output = model(input)

# 导出为ONNX
torch.onnx.export(
    model, 
    (input), 
    "/hy-tmp/result/mobilenetv3-fall.onnx", 
    export_params=True, 
    opset_version=10,  # 设置ONNX的opset版本为10
    do_constant_folding=True, 
    input_names=['input'], 
    output_names=['output']
#     # dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}, 'hidden': {0: 'batch_size'}, 'cell': {0: 'batch_size'}}
)